#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python count_top_blast_organism combined_blastfile.tbl

import sys
import re
from Bio import Entrez

def fmt(f):
    st = '{0:.4}'.format(f)
    return st

Entrez.email = "jimmysaw@gmail.com"

blastfile = sys.argv[1]
bf = open(blastfile, "rU")
bfl = bf.readlines()

si = blastfile.split('.')
seqid = si[0]

hit1 = re.compile('\[(.*)\]')
hit2 = re.compile('\[.*\] .*\[(.*)\]')
hit3 = re.compile('\[.*\] .*\[.*\] .*\[(.*)\]')
hit4 = re.compile('\[.*\]-.* \[(.*)\]')

lin = re.compile('(\w+ \w+).*')
#lin = re.compile('(\w+).*')
acc = re.compile('\w+|(.*)|')

num = len(bfl)

locus_list = []
top_acc_list = []
top_org_list = []

for line in bfl:
    top_acc = ""
    top_org = ""
    l = line.split('\t')
    locus_tag = l[5]
    if acc.match(l[9]):
        p = acc.match(l[9])
        top_acc = p.group(1)
    else:
        top_acc = "No accession"
    desc = l[13].rstrip()
    if hit4.search(desc):
        p = hit4.search(desc)
        top_org = p.group(1)
    elif hit3.search(desc):
        p = hit3.search(desc)
        top_org = p.group(1)
    elif hit2.search(desc):
        p = hit2.search(desc)
        top_org = p.group(1)
    elif hit1.search(desc):
        p = hit1.search(desc)
        top_org = p.group(1)
    else:
        top_org = "No hit"
    locus_list.append(locus_tag)
    top_acc_list.append(top_acc)
    top_org_list.append(top_org)

unique_orgs = set(top_org_list)
unique_org_list = list(sorted(unique_orgs))
org_names = []
org_counts = []
for d in unique_org_list:
    org_names.append(d)
    count = top_org_list.count(d)
    org_counts.append(count)
largest = max(org_counts)
smallest = min(org_counts)
max_index = org_counts.index(largest)
min_index = org_counts.index(smallest)

#print "Top hit organism:" + "\t" + org_names[max_index] + "\t" + str(largest)

for x in unique_org_list:
    print str(top_org_list.count(x)) + "\t" + x

bf.close()
